#include "../inc/ff_utils.h"
#include "../inc/ff_utils_cuda.cuh"

using namespace feifei;


__global__ void vectorAdd(const float *A, const float *B, float *C, int ElementNum)
{
    int i = blockDim.x * blockIdx.x + threadIdx.x;

    if (i < ElementNum)
    {
        C[i] = A[i] + B[i];
    }
}

cudaError_t RtnState;
int ErrorCnt = 0;
void ChkRtn()
{
    if (RtnState != cudaSuccess)
        ErrorCnt++;
}

int ElementNum = 50000;
size_t ElementSize = ElementNum * sizeof(float);
int main(int argc, char* argv[])
{
    PrintSeperator('*');
    PrintHeader("cuda vector_add", '*');
    PrintSeperator('*');
    
    InitGpuRuntime();

    // allocate cpu memory
    float *h_A = (float *)malloc(ElementSize);
    float *h_B = (float *)malloc(ElementSize);
    float *h_C = (float *)malloc(ElementSize);

    // init 
    for (int i = 0; i < ElementNum; ++i)
    {
        h_A[i] = i * 100.0f;
        h_B[i] = i * -0.01f;
    }

    // allocate gpu memory
    LOG("allocate gpu memory");
    float *d_A = NULL;
    RtnState = cudaMalloc((void **)&d_A, ElementSize); ChkRtn();
    float *d_B = NULL;
    RtnState = cudaMalloc((void **)&d_B, ElementSize); ChkRtn();
    float *d_C = NULL;
    RtnState = cudaMalloc((void **)&d_C, ElementSize); ChkRtn();

    // copy to gpu
    LOG("copy to gpu");
    RtnState = cudaMemcpy(d_A, h_A, ElementSize, cudaMemcpyHostToDevice); ChkRtn();
    RtnState = cudaMemcpy(d_B, h_B, ElementSize, cudaMemcpyHostToDevice); ChkRtn();

    // launch kernel
    int threadsPerBlock = 256;
    int blocksPerGrid =(ElementNum + threadsPerBlock - 1) / threadsPerBlock;

    LOG("launch kernel: thread per block = %d", threadsPerBlock);
    LOG("               block   number   = %d", blocksPerGrid);
    LOG("               thread  number   = %d", threadsPerBlock * blocksPerGrid);
    LOG("               element number   = %d", ElementNum);
    vectorAdd<<<blocksPerGrid, threadsPerBlock>>>(d_A, d_B, d_C, ElementNum);

    RtnState = cudaGetLastError(); ChkRtn();

    // copy back
    LOG("copy back");
    RtnState = cudaMemcpy(h_C, d_C, ElementSize, cudaMemcpyDeviceToHost); ChkRtn();

    // verify
    LOG("verify");
    for (int i = 0; i < ElementNum; ++i)
    {
        if (fabs(h_A[i] + h_B[i] - h_C[i]) > 1e-5)
        {
            fprintf(stderr, "Result verification failed at element %d!\n", i);
            exit(EXIT_FAILURE);
        }
    }

    // free gpu
    RtnState = cudaFree(d_A); ChkRtn();
    RtnState = cudaFree(d_B); ChkRtn();
    RtnState = cudaFree(d_C); ChkRtn();

    // free cpu
    free(h_A);
    free(h_B);
    free(h_C);

    LOG("error_cnt = %d", ErrorCnt);
    return ErrorCnt;
}

